/*
 * Copyright (c) 2019, XXX All Rights Reserved.
 */
package net.sf.log4jdbc.agent.driver;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;

/**
 * <p>jdbc Driver 类的二进制代码增强器</p>
 * @author zhanxb
 * @since 2021/8/30
 **/
public class JdbcDriverTransformer extends AbstractTransformer implements ClassFileTransformer
{
    public static final String DRIVER_CONNECT_METHOD = "connect";
    private static final List<String> SUPPORT_DRIVER_CLASSES;
    static
    {
        SUPPORT_DRIVER_CLASSES = new ArrayList<String>();
        // MySQL 5.1.44 之后的驱动
        SUPPORT_DRIVER_CLASSES.add("com.mysql.cj.jdbc.NonRegisteringDriver");
        // MySQL 5.1.44 之前的驱动
        SUPPORT_DRIVER_CLASSES.add("com.mysql.jdbc.NonRegisteringDriver");
        // Oracle的驱动
        SUPPORT_DRIVER_CLASSES.add("oracle.jdbc.driver.OracleDriver");
        // PostgreSQL的驱动
        SUPPORT_DRIVER_CLASSES.add("org.postgresql.Driver");
    }

    /**
     * 判断类是否支持。
     *
     * @param className 类的全类名
     * @return true 表示支持
     */
    public static boolean isSupport(String className)
    {
        if (className == null || className.length() <= 0)
        {
            return false;
        }
        return SUPPORT_DRIVER_CLASSES.contains(className);
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException
    {
        try
        {
            if (className!=null && className.length() > 0)
            {
                final String newClassName = className.replaceAll("/", ".");
                if (isSupport(newClassName))
                {
                    final CtClass ctClass = enhanceConnectMethod(getClassPool(loader), newClassName);
                    if (ctClass!=null)
                    {
                        debug(newClassName + " \u589e\u5f3a\u6210\u529f!!!");
                        return ctClass.toBytecode();
                    }
                }
            }
        } catch (Exception e)
        {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 增强Driver类的connect方法。
     *
     * @return 返回增强后的类定义
     * @throws Exception 增强时抛出的异常
     */
    public static CtClass enhanceConnectMethod(String className) throws Exception
    {
        ClassPool classPool = ClassPool.getDefault();
        CtClass driverClass = classPool.get(className);
        if (driverClass != null)
        {
            CtMethod connectMethod = driverClass.getDeclaredMethod(DRIVER_CONNECT_METHOD);
            connectMethod.addLocalVariable("tstart", CtClass.longType);
            connectMethod.insertBefore("{tstart = System.currentTimeMillis();}");
            connectMethod.insertAfter(
                    "{ net.sf.log4jdbc.sql.jdbcapi.ConnectionSpy cspy = new net.sf.log4jdbc.sql.jdbcapi.ConnectionSpy($_,net.sf.log4jdbc.sql.jdbcapi.DriverSpy.getRdbmsSpecifics($_), System.currentTimeMillis() - tstart,net.sf.log4jdbc.log.SpyLogFactory.getSpyLogDelegator());$_=cspy;}");
        }
        return driverClass;
    }

    /**
     * 增强Driver类的connect方法。
     *
     * @return 返回增强后的类定义
     * @throws Exception 增强时抛出的异常
     */
    public static CtClass enhanceConnectMethod(ClassPool classPool, String className) throws Exception
    {
        CtClass driverClass = classPool.get(className);
        if (driverClass != null)
        {
            CtMethod connectMethod = driverClass.getDeclaredMethod(DRIVER_CONNECT_METHOD);
            CtMethod agentMethod =
                    CtNewMethod.copy(connectMethod, connectMethod.getName() + "$agent", driverClass, null);
            driverClass.addMethod(agentMethod);
            StringBuilder connectMethodSrc = new StringBuilder();
            connectMethodSrc.append("{").append("long tstart = System.currentTimeMillis();");
            connectMethodSrc.append("java.sql.Connection c =").append(connectMethod.getName()).append("$agent($$);");
            connectMethodSrc.append(
                    "if (c == null){throw new java.sql.SQLException(\"invalid or unknown driver url: \" + $1);}");
            connectMethodSrc.append(
                    "net.sf.log4jdbc.sql.jdbcapi.ConnectionSpy cspy = new net.sf.log4jdbc.sql.jdbcapi.ConnectionSpy(c,net.sf.log4jdbc.sql.jdbcapi.DriverSpy.getRdbmsSpecifics(c), System.currentTimeMillis() - tstart,net.sf.log4jdbc.log.SpyLogFactory.getSpyLogDelegator());");
            connectMethodSrc.append("return cspy;").append("}");
            connectMethod.setBody(connectMethodSrc.toString());
        }
        return driverClass;
    }
}
